import torch
import torch.nn as nn
import numpy as np


class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(32, 16),   # 对应输入层x
            nn.BatchNorm1d(16),
            nn.ReLU(inplace=True),
            nn.Linear(16, 32),
            nn.BatchNorm1d(32),
            nn.ReLU(inplace=True),
        )
        self.decoder = nn.Sequential(
            nn.Linear(32, 16),
            nn.BatchNorm1d(16),
            nn.ReLU(inplace=True),
            nn.Linear(16, 5),   # 对应输出层y
        )
    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return decoded

model_path = "./logs/RAdam32_Epoch500.pth"  #需要修改
file_path_cesi = "./data/data_ce32.txt"
model = MLP()
model.eval()
state_dict = torch.load(model_path)
model.load_state_dict(state_dict,strict=False)
model = nn.DataParallel(model)


def load_data(file_path):
    '''导入数据
    input:  file_path(string):文件的存储位置
    output: data(mat):数据
    '''
    f = open(file_path)
    for line in f.readlines():
        lines = line.strip().split("\t")
        for x in lines:
            #print(x)
            rgb = np.fromstring(x, dtype=float, sep=' ')
            rgb_tensor = torch.from_numpy(rgb).float()
            out = model(rgb_tensor.unsqueeze(0))
            # 原始的预测结果，离 正10 最近的为预测类型
            print("%4.2f      %4.2f      %4.2f      %4.2f      %4.2f" % (out[0][0], out[0][1], out[0][2], out[0][3], out[0][4]))
            # 转化之后的预测结果，最小的为预测的类型
            #print("%6.4f      %6.4f      %6.4f      %6.4f      %6.4f" % (abs(out[0][0]-10), abs(out[0][1]-10),abs(out[0][2]-10),abs(out[0][3]-10),abs(out[0][4]-10)))

load_data(file_path_cesi)